/* -*- Mode: C; indent-tabs-mode: t; c-basic-offset: 4; tab-width: 4 -*-  */
/*
 * catchit
 * Copyright (C) 2013 Ryan Grinkewitz <graphite@Toto>
 */

#include "q-learning.h"
#include <iostream>


QLearning::QLearning(int stateSize,int actionSize,
                     float learningRate, int discountFactor){
	this->stateSize = stateSize;
	this->actionSize = actionSize;
	this->learningRate = learningRate;
	this->discountFactor = discountFactor;

	rewards = new float* [stateSize];
	for(int i = 0; i < stateSize; ++i) {
		rewards[i] = new float[actionSize];
		for(int j = 0; j < actionSize; j++){
			rewards[i][j] = 0;
		}
	}

	srand(time(NULL));
}

QLearning::~QLearning(){
	for(int i = 0; i < stateSize; ++i) {
		delete rewards[i];
	}

	delete rewards;
}

void QLearning::setCurrentState(int state){
	this->currentState = state;
}

//e-greedy
int QLearning::getAction(){
	float bestReward = rewards[currentState][0];
	int bestAction = 0;
	bool allEqual = false;
	if(bestReward == 0){
		allEqual = true;
	}
	
	for(int i = 1; i < actionSize; i++){
		float nextReward = rewards[currentState][i];
		if(bestReward < nextReward){
			bestReward = nextReward;
			bestAction = i;
			allEqual = false;
		}
	}

	if(rand() % 10 == 1 || allEqual){
		bestAction = rand() % actionSize;
		std::cout << "***Generating random action bestReward = "<< bestReward << std::endl;
	}

	currentAction = bestAction;
	return bestAction;
	
}

void QLearning::updateQFunction(int newState, float reward){
	float qFunction =  rewards[currentState][currentAction];


	float bestNewStateReward = rewards[newState][0];
	for(int i = 1; i < actionSize; i++){
		int nextReward = rewards[newState][i];
		if(bestNewStateReward < nextReward){
			bestNewStateReward = nextReward;
		}
	}

	qFunction = qFunction + learningRate * 
		(reward + (discountFactor * bestNewStateReward) - qFunction);
	
	rewards[currentState][currentAction] = qFunction;

	//std::cout << "qFunction:" << qFunction << std::endl;
	
}

void QLearning::readQFunctionTable(std::string qFunctionTable){
	std::ifstream fileStream;
	std::string fileLine;
	std::string fileColumn;
	fileStream.open(qFunctionTable.c_str());

	int action = 0;
	int state = 0;		
	
	while(!fileStream.eof()){
		std::getline(fileStream, fileLine);
		std::istringstream fileLineStream(fileLine);
		std::getline(fileLineStream, fileColumn, ',');
		action = atoi(fileColumn.c_str());
		std::getline(fileLineStream, fileColumn, ',');
		state = atoi(fileColumn.c_str());
		std::getline(fileLineStream, fileColumn);
		rewards[state][action] = atof(fileColumn.c_str());
			
	}
}

void QLearning::writeQFunctionTable(std::string qFunctionTable){
	std::ofstream fileStream;
	fileStream.open(qFunctionTable.c_str());
	
	for(int i = 0; i < stateSize; ++i) {
		for(int j = 0; j < actionSize; j++){
			fileStream << j << "," << i << "," << rewards[i][j] << "\n";
		}
	}

	fileStream.close();
}

